from fastapi import FastAPI
from fastapi import WebSocket
from fastapi.templating import Jinja2Templates
from starlette import websockets

from . import login
from . import session_save
from .userobj import User

import asyncio

ws_tasks = []


def main(app: FastAPI, templates: Jinja2Templates):
    @app.websocket("/api/home/ws")
    @login.check_cookie
    async def ws_connect(user: User, websocket: WebSocket):  # 监听桌面 ws 请求
        create_tasks = []
        try:
            await websocket.accept()
            for i in ws_tasks:  # 执行任务
                create_tasks.append(asyncio.create_task(
                    i(websocket)))
            while True:
                data = await websocket.receive_json()
                if (data["type"] == "save_session"):  # 保存会话
                    await session_save.set_session_window(user, data["data"])
        except websockets.WebSocketDisconnect:
            for i in create_tasks:  # 关闭任务
                i.cancel()
